"""
处理排序语料
"""
import os
from tqdm import tqdm

from lib import cut

# 获取问题语料路径
q_path = r"F:\virtual_environment\AI_Study\AI_study_code\人工智能NLP项目\案例-chat_service" \
         r"\corpus\dnn\sort/merged_q.txt"
# 获取相似问题语料路径
sim_q_path = r"F:\virtual_environment\AI_Study\AI_study_code\人工智能NLP项目\案例-chat_service" \
             r"\corpus\dnn\sort/merged_sim_q.txt"
# 获取判断是否正确语料路径
v_path = r"F:\virtual_environment\AI_Study\AI_study_code\人工智能NLP项目\案例-chat_service" \
         r"\corpus\dnn\sort/merged_v.txt"


def process_sort_corpus(by_word=False):
    # 获取数据
    f_q = open(q_path, mode="r", encoding="UTF-8")
    f_sim_q = open(sim_q_path, mode="r", encoding="UTF-8")
    f_v = open(v_path, mode='r', encoding="UTF-8")

    q_data = f_q.readlines()
    sim_q_data = f_sim_q.readlines()
    v_data = f_v.readlines()

    f_q.close()
    f_sim_q.close()
    f_v.close()

    save_q_file = open(os.path.join(os.path.dirname(q_path), "q_cut_by_word.txt" if by_word else "q_cut.txt"),
                       mode="a", encoding="UTF-8")
    save_sim_q_file = open(
        os.path.join(os.path.dirname(sim_q_path), "sim_q_cut_by_word.txt" if by_word else "sim_q_cut.txt"), mode="a",
        encoding="UTF-8")
    save_v_file = open(os.path.join(os.path.dirname(v_path), "target.txt"), mode='a', encoding="UTF-8")

    # 进行切分并保存数据
    for q_line, sim_q_line, v_line in tqdm(zip(q_data[::25], sim_q_data[::25], v_data[::25]), ascii=True,
                                           desc="按单个字对问题进行切分" if by_word else "按词对问题进行切分"):
        q_line = " ".join(cut(q_line.strip(), by_word=by_word)) + '\n'
        save_q_file.write(q_line)
        sim_q_line = " ".join(cut(sim_q_line.strip(), by_word=by_word)) + '\n'
        save_sim_q_file.write(sim_q_line)
        save_v_file.write(v_line.strip() + '\n')

    save_q_file.close()
    save_sim_q_file.close()
    save_v_file.close()


if __name__ == '__main__':
    process_sort_corpus()
